import numpy as np
import pyrealsense2 as rs
from ultralytics import YOLO
import cv2
import time
class YOLO11:
    def __init__(self, model_path):
        self.model = YOLO(model_path)
        self.class_dict={0: 'high', 1: 'low', 2: 'normal'}
        self.class_dict_predict = {'high': 0, 'low':0, 'normal':0}
        self.nowtime = time.time()
        self.lasttime = time.time() 
    def detect_objects(self, image_path):
        # Load the image
        image = cv2.imread(image_path)
        
        # Perform object detection
        results = self.model.predict(source=image, conf=0.5, show=True)
        
        # Process results
        for result in results:
            boxes = result.boxes
            for box in boxes:
                x1, y1, x2, y2 = box.xyxy[0]
                conf = box.conf[0]
                cls = box.cls[0]
                print(f"Detected {cls} with confidence {conf} at [{x1}, {y1}, {x2}, {y2}]")
    def predict(self, image,IsStream=False):
            results = self.model.predict(source=image,streem=IsStream)
            return results
    def analyze(self,results):
        for result in results:
            prob_top1 = result.probs.top1conf.item()
            # 筛选
            self.nowtime = time.time()
            if prob_top1 < 0.9:
                continue
            num_class_top1 = result.probs.top1
            class_top1 = self.class_dict[num_class_top1]
            self.predict_list.append((class_top1, prob_top1))
            self.lasttime = self.nowtime
            
        if len(self.predict_list) >= 10:
            # 统计
            for item in self.predict_list:
                self.class_dict_predict[item[0]] += 1
            max_class = max(self.class_dict_predict, key=self.class_dict_predict.get)
            # 清空计数 
            for name, value in self.class_dict_predict.items():
                self.class_dict_predict[name] = 0
            # 清空预测列表
            self.predict_list.clear()
            # 打印结果
            print("the class is",max_class)
    def timeout_clear(self):
        if self.nowtime - self.lasttime > 2:
            # 清空计数 
            for name, value in self.class_dict_predict.items():
                self.class_dict_predict[name] = 0
            # 清空预测列表
            self.predict_list.clear()
            self.lasttime = self.nowtime
    def task_predict(self, image,IsStream=False):
        self.analyze(self.predict(image,IsStream))
        self.timeout_clear()
    def test(self,image):
        results=self.predict(image)
        for result in results:
            prob_top1 = result.probs.top1conf.item()
            # 筛选
            self.nowtime = time.time()
            if prob_top1 < 0.9:
                continue
            num_class_top1 = result.probs.top1
            class_top1 = self.class_dict[num_class_top1]
            self.predict_list.append((class_top1, prob_top1))
            self.lasttime = self.nowtime
            print("类别是",class_top1)
class Rs2camera:
    def __init__(self, width=640, height=480, fps=30):
        self.pipeline = rs.pipeline()
        self.config = rs.config()
        self.pipeline_wrapper = rs.pipeline_wrapper(self.pipeline)
        self.pipeline_profile = self.config.resolve(self.pipeline_wrapper)
        self.device = self.pipeline_profile.get_device()
        self.device_product_line = str(self.device.get_info(rs.camera_info.product_line))
        found_rgb = False
        for s in self.device.sensors:
            if s.get_info(rs.camera_info.name) == 'RGB Camera':
                found_rgb = True
                break
        if not found_rgb:
            print("The demo requires Depth camera with Color sensor")
            exit(0)
        self.config.enable_stream(rs.stream.color, width, height, rs.format.bgr8, fps)
        self.config.enable_stream(rs.stream.depth, width, height, rs.format.z16, fps)   
        self.pipeline.start(self.config)    
        self.depth_sensor = self.profile.get_device().first_depth_sensor()
        self.depth_scale = self.depth_sensor.get_depth_scale()
        self.align_to = rs.stream.color
        self.align = rs.align(self.align_to)
    def get_frame(self):
        self.frames = self.pipeline.wait_for_frames()
        self.aligned_frames = self.align.process(self.frames)
        self.color_frame = self.aligned_frames.get_color_frame()
        self.depth_frame = self.aligned_frames.get_depth_frame()
        if not self.color_frame or not self.depth_frame:
            print("No frames received")
            return None, None
        self.color_np= np.asanyarray(self.color_frame.get_data())
        self.depth_np = np.asanyarray(self.depth_frame.get_data())  
        return self.color_np, self.depth_np


        
        
            

        